import json

import pandas as pd
import torch
import onnxruntime as ort
import numpy as np
from matplotlib import pyplot as plt

from model.model_pytorch import Net

class Config:
    # 数据参数
    feature_columns = list(range(1, 10))  # 要作为feature的列，按原数据从0开始计算，也可以用list 如 [2,4,6,8] 设置
    label_columns = [9]  # 要预测的列，按原数据从0开始计算, 如同时预测第四，五列 最低价和最高价
    # label_in_feature_index = [feature_columns.index(i) for i in label_columns]  # 这样写不行
    label_in_feature_index = (lambda x, y: [x.index(i) for i in y])(feature_columns, label_columns)  # 因为feature不一定从0开始

    predict_day = 72  # 预测未来几项

    # 网络参数
    input_size = len(feature_columns)
    output_size = len(label_columns)

    hidden_size = 128  # LSTM的隐藏层大小，也是输出大小
    lstm_layers = 4  # LSTM的堆叠层数
    dropout_rate = 0.2  # dropout概率
    time_step = 168  # 设置用前多少项的数据来预测，也是LSTM的time step数，请保证训练数据量大于它
    time_Interval_step = 4  # 训练数据使用4小时间隔：1-20 5-24.。。

    # 训练参数
    use_cuda = False  # 是否使用GPU训练

    train_data_rate = 0.98  # 训练数据占总体数据比例，测试数据就是 1-train_data_rate
    valid_data_rate = 0.15  # 验证数据占训练数据比例，验证集在训练过程使用，为了做模型和参数选择

    learning_rate = 0.001

    # 训练模式
    debug_mode = False  # 调试模式下，是为了跑通代码，追求快
    debug_num = 500  # 仅用debug_num条数据来调试

    # 框架参数
    model_postfix = {"pytorch": ".pth", "keras": ".h5", "tensorflow": ".ckpt"}

    # 路径参数
    predict_data_path = "../data/TrainDataset.csv"

con = Config()

# # 加载模型
# model = Net(con)
# model.load_state_dict(torch.load('../checkpoint/pytorch/model_pytorch.pth'))
#
# # # 保存为ONNX模型
# input_shape = (1, con.time_step, 9)
input_names = ["input"]
output_names = ["output"]
# x = torch.randn(input_shape)
# torch.onnx.export(model, x, "../checkpoint/pytorch/lstm_model.onnx", input_names=input_names, output_names=output_names)

# 加载ONNX模型
ort_session = ort.InferenceSession("../Flask/data/lstm_model.onnx")

# 从 JSON 文件中读取平均数和方差
with open("../checkpoint/pytorch/config.json", "r") as f:
    config = json.load(f)
mean = np.array(config["mean"])
std = np.array(config["std"])

# 预测未来的值
future = 72
init_data = pd.read_csv(con.predict_data_path, usecols=con.feature_columns).values

x_test = init_data[-con.time_step:]

x_test = (x_test -mean)/std

# predictions = np.array(object=object)
with torch.no_grad():
    # for i in range(future):
        x_input = np.array(x_test,dtype=np.float32).reshape((1,con.time_step,9))
        ort_inputs = {input_names[0]: x_input}
        ort_outs = ort_session.run(output_names, ort_inputs)
        prediction = ort_outs[0].reshape(-1)
        prediction = prediction * std[8] + mean[8]
        print("predictions-",prediction)
        # predictions.append(prediction)
        # np.append(predictions,prediction)
        # x_test = np.concatenate((x_test, prediction))

# print("predictions-",predictions)
# 绘制预测结果
# plt.figure(figsize=(12, 6))
# plt.plot(np.arange(len(data)), data, label='Original Data')
# plt.plot(np.arange(len(data), len(data) + future), predictions, label='Predictions')
# plt.legend()
# plt.show()
